#pragma once
#include <bits/stdc++.h>

using namespace std;

enum Color
{
    RED,
    BLACK
};

template <typename V>
struct RBTreeNode
{
    RBTreeNode(const V &data)
        : _left(nullptr),
          _right(nullptr),
          _parent(nullptr),
          _data(data),
          _col(RED)
    {
    }
    /* data */

    struct RBTreeNode<V> *_left;
    struct RBTreeNode<V> *_right;
    struct RBTreeNode<V> *_parent;
    Color _col;
    V _data;
};
template <typename K, typename V, typename keyOfValue>  
class RBTree
{
    typedef RBTreeNode<V> Node;
private:
    keyOfValue val;

public:
    bool insert(const V &data)
    {
        //默认根节点为黑色
        if (_root == nullptr)
        {
            _root = new Node(data);
            _root->_col = BLACK;
            return true;
        }
        Node *cur = _root;
        Node *parent = nullptr;
        while (cur)
        {
            if (val(cur->_data) < val(data))
            {
                parent = cur;
                cur = cur->_right;
            }
            else if (val(cur->_data) > val(data))
            {
                parent = cur;
                cur = cur->_left;
            }
            else
            {
                return false;
            }
        }
        //插入结点在左侧或右侧（默认是红色结点）
        cur = new Node(data);
        if (val(data) > val(parent->_data))
        {
            parent->_right = cur;
        }
        else
        {
            parent->_left = cur;
        }
        parent = cur->_parent;
        while (parent && parent->_col == RED)
        {
            Node *grandParent = parent->_parent;
            if (parent == grandParent->_left)
            {
                Node *uncle = grandParent->_right;
                //情况一：uncle存在，则只需要变色即可
                if (uncle && uncle->_col == RED)
                {
                    parent->_col = BLACK;
                    uncle->_col = BLACK;
                    grandParent->_col = RED;

                    cur = grandParent;
                    parent = cur->_parent;
                }
                //情况二、三：uncle不存在或者uncle为黑色
                else
                {
                    //情况二，cur，p和g结点在同一边
                    if (cur == parent->_left)
                    {
                        rotateR(grandParent);
                        parent->_col = BLACK;
                        grandParent->_col = RED;
                    }
                    else
                    {
                        //情况三，cur，p和g结点不在同一边
                        rotateL(parent);
                        rotateR(grandParent);
                        cur->_col = BLACK;
                        grandParent->_col = RED;
                    }
                    break;
                }
            }
            else
            {
                Node *uncle = grandParent->_left;
                //情况一：uncle存在，则只需要变色即可
                if (uncle && uncle->_col == RED)
                {
                    parent->_col = BLACK;
                    uncle->_col = BLACK;
                    grandParent->_col = RED;

                    cur = grandParent;
                    parent = cur->_parent;
                }
                //情况二、三：uncle不存在或者uncle为黑色
                else
                {
                    //情况二，cur，p和g结点在同一边
                    if (cur == parent->_left)
                    {
                        rotateL(grandParent);
                        parent->_col = BLACK;
                        grandParent->_col = RED;
                    }
                    else
                    {
                        //情况三，cur，p和g结点不在同一边
                        rotateR(parent);
                        rotateL(grandParent);
                        cur->_col = BLACK;
                        grandParent->_col = RED;
                    }
                    break;
                }
            }
        }
        _root->_col = BLACK;
        return true;
    }
    void InOrder()
    {
        _InOrder(_root);
        cout << endl;
    }
    bool isBlance()
    {
        return _isBlance(_root);
    }

private:
    bool PrevCheck(Node *root, int blackNum, int &benchmark)
    {
        if (root == nullptr)
        {
            if (benchmark == 0)
            {
                benchmark = blackNum;
                return true;
            }

            if (blackNum != benchmark)
            {
                cout << "某条黑色节点的数量不相等" << endl;
                return false;
            }
            else
            {
                return true;
            }
        }

        if (root->_col == BLACK)
        {
            ++blackNum;
        }

        if (root->_col == RED && root->_parent && root->_parent->_col == RED)
        {
            cout << "存在连续的红色节点" << endl;
            return false;
        }

        return PrevCheck(root->_left, blackNum, benchmark) && PrevCheck(root->_right, blackNum, benchmark);
    }

    bool _isBlance(Node *root)
    {
        if (root == nullptr)
        {
            return true;
        }
        if (_root->_col != BLACK)
        {
            cout << "根节点不是黑色" << endl;
            return false;
        }

        int benchMark = 0;
        return PrevCheck(_root, 0, benchMark);
    }
    int height(Node *root)
    {
        if (root == nullptr)
        {
            return 0;
        }
        return 1 + max(height(root->_left), height(root->_right));
    }
    void _InOrder(Node *root)
    {
        if (root == nullptr)
        {
            return;
        }
        _InOrder(root->_left);
        std::cout << root->_data.first << " " << root->_data.second << endl;
        _InOrder(root->_right);
    }
    void rotateL(Node *root)
    {
        Node *rootR = root->_right;
        Node *rootRL = rootR->_left;

        root->_right = rootRL;
        if (rootRL)
        {
            rootRL->_parent = root;
        }
        Node *rootParent = root->_parent;

        rootR->_left = root;
        root->_parent = rootR;
        if (root == _root)
        {
            _root = rootR;
            rootR->_parent = nullptr;
        }
        else
        {
            if (root == rootParent->_left)
            {
                rootParent->_left = rootR;
            }
            else
            {
                rootParent->_right = rootR;
            }
            rootR->_parent = rootParent;
        }

    }
    void rotateR(Node *root)
    {
        Node *rootL = root->_left;
        Node *rootLR = rootL->_right;

        root->_left = rootLR;
        if (rootLR)
        {
            rootLR->_parent = root;
        }
        Node *rootParent = root->_parent;

        rootL->_right = root;
        root->_parent = rootL;
        if (root == _root)
        {
            _root = rootL;
            rootL->_parent = nullptr;
        }
        else
        {
            if (rootParent->_left == root)
            {
                rootParent->_left = rootL;
            }
            else
            {
                rootParent->_right = rootL;
            }
            rootL->_parent = rootParent;
        }
    }

private:
    Node *_root = nullptr;
};